import os
import json
import logging
import pandas as pd
import re
from typing import Any, Dict, List, Tuple, Optional
from dataclasses import dataclass
from datetime import datetime

# langchain
from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate
from langchain_ollama import OllamaLLM, OllamaEmbeddings

# llama-index & chroma
import chromadb
from llama_index.core import Settings
from llama_index.core import Document
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.postprocessor import SimilarityPostprocessor

# 日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class LogMetadata:
    """日志元数据结构"""
    service: str
    level: str
    error_code: str
    component: str
    cause: str
    timestamp: Optional[str] = None

class EnhancedRAGSystem:
    def __init__(
        self,
        log_path: str,
        knowledge_path: str,
        llm: str = "deepseek-r1:7b",
        embedding_model: str = "bge-large:latest"
    ) -> None:
        # 初始化模型
        self.embedding_model = OllamaEmbeddings(model=embedding_model)
        self.llm = OllamaLLM(model=llm, temperature=0.1)

        # 设置全局配置
        Settings.llm = self.llm
        Settings.embed_model = self.embedding_model

        # 路径配置
        self.log_path = log_path
        self.knowledge_path = knowledge_path

        # 存储索引
        self.log_index = None
        self.knowledge_index = None
        self.vector_store = None

        # 构建向量存储
        self._build_enhanced_vectorstore()

    def _build_enhanced_vectorstore(self):
        """构建增强的向量存储系统"""
        vector_store_path = "./data/vector_stores"
        os.makedirs(vector_store_path, exist_ok=True)

        chroma_client = chromadb.PersistentClient(path=vector_store_path)

        # 创建日志集合
        log_collection = chroma_client.get_or_create_collection("enhanced_log_collection")
        log_vector_store = ChromaVectorStore(chroma_collection=log_collection)
        log_storage_context = StorageContext.from_defaults(vector_store=log_vector_store)

        # 构建日志索引
        if log_documents := self._load_enhanced_log_documents(self.log_path):
            self.log_index = VectorStoreIndex.from_documents(
                log_documents,
                storage_context=log_storage_context,
                show_progress=True,
            )
            logger.info(f"增强日志库索引构建完成，共 {len(log_documents)} 条日志")

        # 创建专业知识集合
        knowledge_collection = chroma_client.get_or_create_collection("knowledge_collection")
        knowledge_vector_store = ChromaVectorStore(chroma_collection=knowledge_collection)
        knowledge_storage_context = StorageContext.from_defaults(vector_store=knowledge_vector_store)

        # 构建知识库索引
        if knowledge_documents := self._load_knowledge_documents(self.knowledge_path):
            self.knowledge_index = VectorStoreIndex.from_documents(
                knowledge_documents,
                storage_context=knowledge_storage_context,
                show_progress=True,
            )
            logger.info(f"专业知识库索引构建完成，共 {len(knowledge_documents)} 条知识")

    def _parse_log_metadata(self, log_content: str) -> LogMetadata:
        """解析日志元数据"""
        # 尝试从CSV格式解析
        # CSV格式：服务,级别,错误,消息,组件,原因
        if ',' in log_content:
            parts = log_content.split(',')
            if len(parts) >= 6:
                return LogMetadata(
                    service=parts[0].strip() if parts[0] else "",
                    level=parts[1].strip() if parts[1] else "",
                    error_code=parts[2].strip() if parts[2] else "",
                    component=parts[4].strip() if parts[4] else "",  # 组件是第5列(索引4)
                    cause=parts[5].strip() if parts[5] else ""
                )

        # 从正则表达式解析
        patterns = {
            'service': r'(?:服务|Service)[：:]\s*(\w+)',
            'level': r'(?:级别|Level)[：:]\s*(\w+)',
            'error': r'(?:错误|Error)[：:]\s*(\w+)',
            'component': r'(?:组件|Component)[：:]\s*(\w+)'
        }

        metadata = LogMetadata("", "", "", "", "")
        for key, pattern in patterns.items():
            match = re.search(pattern, log_content, re.IGNORECASE)
            if match:
                if key == 'service':
                    metadata.service = match.group(1)
                elif key == 'level':
                    metadata.level = match.group(1)
                elif key == 'error':
                    metadata.error_code = match.group(1)
                elif key == 'component':
                    metadata.component = match.group(1)

        return metadata

    def _load_enhanced_log_documents(self, data_path: str) -> List[Document]:
        """加载增强的日志文档"""
        if not os.path.exists(data_path):
            logger.warning(f"数据路径不存在: {data_path}")
            return []

        documents = []
        for file in os.listdir(data_path):
            ext = os.path.splitext(file)[1]
            if ext not in [".txt", ".md", ".json", ".jsonl", ".csv"]:
                continue

            file_path = f"{data_path}/{file}"
            try:
                if ext == ".csv":
                    chunk_size = 1000
                    for chunk in pd.read_csv(file_path, chunksize=chunk_size):
                        for idx, row in chunk.iterrows():
                            # 直接使用DataFrame的列名访问，更可靠
                            try:
                                service = str(row.get('服务', '')).strip() if pd.notna(row.get('服务')) else ""
                                level = str(row.get('级别', '')).strip() if pd.notna(row.get('级别')) else ""
                                error_code = str(row.get('错误', '')).strip() if pd.notna(row.get('错误')) else ""
                                message = str(row.get('消息', '')).strip() if pd.notna(row.get('消息')) else ""
                                component = str(row.get('组件', '')).strip() if pd.notna(row.get('组件')) else ""
                                cause = str(row.get('原因', '')).strip() if pd.notna(row.get('原因')) else ""
                            except KeyError:
                                # 如果列名不匹配，fallback到原来的方法
                                content = str(row).replace("Pandas", " ")
                                metadata = self._parse_log_metadata(content)
                                structured_content = self._structure_log_content(content, metadata)
                                doc_metadata = {
                                    "service": metadata.service,
                                    "level": metadata.level,
                                    "error_code": metadata.error_code,
                                    "component": metadata.component,
                                    "document_type": "log"
                                }
                                document = Document(text=structured_content, metadata=doc_metadata)
                                documents.append(document)
                                continue
                            
                            # 构建完整的日志内容
                            content = f"{service},{level},{error_code},{message},{component},{cause}"
                            
                            # 创建元数据对象
                            metadata = LogMetadata(
                                service=service,
                                level=level,
                                error_code=error_code,
                                component=component,
                                cause=cause
                            )

                            # 创建结构化文档
                            structured_content = self._structure_log_content(content, metadata)

                            # 添加元数据
                            doc_metadata = {
                                "service": metadata.service,
                                "level": metadata.level,
                                "error_code": metadata.error_code,
                                "component": metadata.component,
                                "document_type": "log"
                            }

                            document = Document(
                                text=structured_content,
                                metadata=doc_metadata
                            )
                            documents.append(document)
                else:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        content = f.read()
                        metadata = self._parse_log_metadata(content)
                        structured_content = self._structure_log_content(content, metadata)

                        doc_metadata = {
                            "service": metadata.service,
                            "level": metadata.level,
                            "error_code": metadata.error_code,
                            "component": metadata.component,
                            "document_type": "log"
                        }

                        document = Document(
                            text=structured_content,
                            metadata=doc_metadata
                        )
                        documents.append(document)
            except Exception as e:
                logger.error(f"加载增强文档失败 {file_path}: {e}")

        return documents

    def _structure_log_content(self, content: str, metadata: LogMetadata) -> str:
        """结构化日志内容以提高检索效果"""
        structured = f"""
故障事件描述：
{content}

结构化信息：
- 服务组件：{metadata.service}
- 日志级别：{metadata.level}
- 错误代码：{metadata.error_code}
- 故障组件：{metadata.component}
- 故障原因：{metadata.cause}

关键词提取：
{self._extract_keywords(content)}
        """.strip()
        return structured

    def _extract_keywords(self, content: str) -> str:
        """提取关键词"""
        # 常见的故障相关关键词
        fault_keywords = [
            "连接", "超时", "失败", "错误", "异常", "崩溃", "内存", "CPU", "磁盘",
            "网络", "数据库", "认证", "授权", "配置", "依赖", "版本", "兼容",
            "性能", "泄漏", "阻塞", "死锁", "熔断", "降级", "限流", "缓存"
        ]

        keywords = []
        content_lower = content.lower()

        for keyword in fault_keywords:
            if keyword in content_lower:
                keywords.append(keyword)

        return ", ".join(keywords) if keywords else "无特定关键词"

    def _load_knowledge_documents(self, knowledge_path: str) -> List[Document]:
        """加载专业知识库文档"""
        if not os.path.exists(knowledge_path):
            logger.info(f"知识库路径不存在，创建默认知识库: {knowledge_path}")
            self._create_default_knowledge_base(knowledge_path)

        documents = []
        try:
            for file in os.listdir(knowledge_path):
                if file.endswith('.txt'):
                    file_path = f"{knowledge_path}/{file}"
                    with open(file_path, 'r', encoding='utf-8') as f:
                        content = f.read()

                        document = Document(
                            text=content,
                            metadata={
                                "document_type": "knowledge",
                                "source_file": file
                            }
                        )
                        documents.append(document)
        except Exception as e:
            logger.error(f"加载知识库失败: {e}")

        return documents

    def _create_default_knowledge_base(self, knowledge_path: str):
        """创建默认的专业知识库"""
        os.makedirs(knowledge_path, exist_ok=True)

        # 创建故障诊断知识库
        fault_diagnosis_knowledge = """
# 微服务故障诊断知识库

## 1. 数据库连接问题

### 问题特征：
- 错误信息包含"连接池耗尽"、"连接超时"、"数据库连接失败"
- 相关组件：DBConnector、DataSource、ConnectionPool

### 常见原因：
1. 连接泄漏：应用程序未正确释放数据库连接
2. 连接池配置不当：最大连接数设置过小
3. 数据库性能问题：查询响应慢导致连接占用时间长
4. 网络问题：应用与数据库之间的网络延迟或中断

### 解决方案：
1. 检查连接池配置，适当增加最大连接数
2. 添加连接监控和自动回收机制
3. 优化慢查询，添加索引
4. 设置连接超时和空闲连接回收策略
5. 使用连接池监控工具

## 2. 内存问题诊断

### 问题特征：
- 错误信息包含"OutOfMemoryError"、"内存不足"、"内存泄漏"
- 相关组件：JVM、.NET Runtime、Container

### 常见原因：
1. 内存泄漏：对象未及时释放
2. 内存配置不足：JVM堆内存设置过小
3. 大对象分配：单次分配过大内存
4. 缓存数据过多：缓存策略不合理

### 解决方案：
1. 使用内存分析工具（MAT、VisualVM）分析内存使用
2. 调整JVM参数：-Xms、-Xmx、-XX:MaxMetaspaceSize
3. 实施缓存淘汰策略
4. 添加内存监控和告警
5. 考虑使用内存优化技术和数据结构

## 3. 网络与服务调用问题

### 问题特征：
- 错误信息包含"超时"、"连接拒绝"、"网络不可达"
- 相关组件：HttpClient、Feign、RestTemplate

### 常见原因：
1. 目标服务不可用或响应慢
2. 网络延迟或丢包
3. 负载均衡配置问题
4. 熔断降级策略配置不当

### 解决方案：
1. 实施服务健康检查
2. 配置合理的超时时间和重试策略
3. 使用熔断器模式（Hystrix、Resilience4j）
4. 添加服务监控和链路追踪
5. 考虑服务降级和异步处理

## 4. 认证授权问题

### 问题特征：
- 错误信息包含"Token校验失败"、"认证失败"、"权限不足"
- 相关组件：AuthService、JWT、OAuth2

### 常见原因：
1. Token过期或格式错误
2. 签名算法不匹配
3. 密钥配置错误
4. 权限配置不当

### 解决方案：
1. 检查Token生成和验证逻辑
2. 统一签名算法配置
3. 实施Token自动刷新机制
4. 添加详细的认证日志
5. 定期轮换密钥

## 5. 系统资源问题

### 问题特征：
- 错误信息包含"CPU使用率高"、"磁盘空间不足"、"文件描述符耗尽"
- 相关组件：SystemMonitor、OS、Container

### 常见原因：
1. 进程数量过多
2. 文件句柄泄漏
3. 磁盘空间不足
4. CPU密集型操作

### 解决方案：
1. 添加资源监控和告警
2. 实施资源限制和隔离
3. 优化代码逻辑，减少资源消耗
4. 定期清理临时文件和日志
5. 考虑水平扩展和负载均衡
        """

        # 写入知识库文件
        with open(f"{knowledge_path}/故障诊断知识库.md", 'w', encoding='utf-8') as f:
            f.write(fault_diagnosis_knowledge)

        # 创建最佳实践指南
        best_practices = """
# 系统运维最佳实践

## 监控告警
1. 建立完善的监控体系
2. 设置合理的告警阈值
3. 实施分级告警机制

## 日志管理
1. 统一日志格式和级别
2. 实施日志聚合和分析
3. 设置日志轮转策略

## 故障处理
1. 建立故障响应流程
2. 准备故障处理手册
3. 定期进行故障演练

## 性能优化
1. 定期进行性能基准测试
2. 建立性能基线
3. 实施渐进式优化
        """

        with open(f"{knowledge_path}/运维最佳实践.md", 'w', encoding='utf-8') as f:
            f.write(best_practices)

        logger.info("默认知识库创建完成")

    def hybrid_retrieve(self, query: str, top_k: int = 10) -> Dict[str, List[Dict]]:
        """混合检索：结合语义检索和元数据过滤"""
        results = {
            "logs": [],
            "knowledge": []
        }

        # 1. 语义检索日志
        if self.log_index:
            log_retriever = self.log_index.as_retriever(similarity_top_k=top_k)
            log_results = log_retriever.retrieve(query)

            for result in log_results:
                results["logs"].append({
                    "content": result.text,
                    "score": result.score,
                    "metadata": result.metadata or {}
                })

        # 2. 语义检索知识库
        if self.knowledge_index:
            knowledge_retriever = self.knowledge_index.as_retriever(similarity_top_k=5)
            knowledge_results = knowledge_retriever.retrieve(query)

            for result in knowledge_results:
                results["knowledge"].append({
                    "content": result.text,
                    "score": result.score,
                    "metadata": result.metadata or {}
                })

        # 3. 重排序和过滤
        results["logs"] = self._rerank_results(query, results["logs"])
        results["knowledge"] = self._rerank_results(query, results["knowledge"])

        return results

    def _rerank_results(self, query: str, results: List[Dict]) -> List[Dict]:
        """重排序检索结果"""
        if not results:
            return results

        # 基于多个维度进行重排序
        for result in results:
            # 基础相似度分数
            base_score = result["score"]

            # 元数据匹配加分
            metadata_bonus = 0
            metadata = result.get("metadata", {})

            # 关键词匹配加分
            query_lower = query.lower()
            content_lower = result["content"].lower()

            # 如果查询中包含特定服务名，且结果也包含，加分
            for service in ["database", "auth", "payment", "log", "cache"]:
                if service in query_lower and service in content_lower:
                    metadata_bonus += 0.1

            # 错误级别匹配
            level = metadata.get("level")
            if level and level in query_lower:
                metadata_bonus += 0.1

            # 计算最终分数
            result["final_score"] = base_score + metadata_bonus

        # 按最终分数排序
        results.sort(key=lambda x: x.get("final_score", 0), reverse=True)

        return results

    def generate_expert_response(self, query: str, context: Dict) -> str:
        """生成专家级响应"""
        prompt = self._build_expert_prompt(query, context)

        try:
            response = self.llm.invoke(prompt)
            return response
        except Exception as e:
            logger.error(f"LLM调用失败: {e}")
            return f"生成响应时出错: {str(e)}"

    def _build_expert_prompt(self, query: str, context: Dict) -> str:
        """构建专家级诊断提示词"""

        # 构建日志上下文
        log_context = ""
        if context.get("logs"):
            log_context = "## 相关历史日志分析:\n"
            for i, log in enumerate(context["logs"][:5], 1):  # 只取前5个最相关的
                log_context += f"""
### 日志案例 {i} (相关度: {log.get('final_score', log.get('score', 0)):.3f})
{log['content']}

"""

        # 构建知识库上下文
        knowledge_context = ""
        if context.get("knowledge"):
            knowledge_context = "## 专业知识参考:\n"
            for i, knowledge in enumerate(context["knowledge"][:3], 1):
                knowledge_context += f"""
### 知识条目 {i} (相关度: {knowledge.get('final_score', knowledge.get('score', 0)):.3f})
{knowledge['content']}

"""

        # 专家级系统提示
        system_prompt = """你是一位资深的系统运维专家和故障诊断工程师，具有以下专业能力：

1. **深厚的技术背景**：精通分布式系统、微服务架构、数据库优化、网络编程等
2. **丰富的故障诊断经验**：处理过数千起生产环境故障，擅长从日志中发现根本原因
3. **系统化分析方法**：采用结构化思维，从现象到原因，再到解决方案
4. **实用的解决方案**：提供可操作、可落地的解决建议

在分析问题时，请遵循以下原则：
- 基于提供的日志和知识进行推理，不要编造信息
- 从多个角度分析问题可能的原因
- 提供具体的操作步骤和验证方法
- 考虑问题的紧急程度和影响范围
- 建议预防措施和长期优化方案

请以专业、准确、实用的方式回答用户问题。"""

        user_prompt = f"""
{log_context}

{knowledge_context}

## 当前需要诊断的问题:
{query}

请基于以上信息，提供专业的故障诊断报告，包括：

1. **问题现象分析**
   - 基于日志描述的关键现象
   - 问题的紧急程度和影响范围评估

2. **可能原因分析**
   - 列出2-3个最可能的原因
   - 每个原因的支持证据和可能性评估

3. **解决方案建议**
   - 立即处理措施（紧急修复）
   - 根本解决方案（长期优化）
   - 具体的操作步骤和验证方法

4. **预防措施**
   - 如何避免类似问题再次发生
   - 监控和告警建议

请确保回答专业、准确、具有可操作性。
"""

        return f"{system_prompt}\n\n{user_prompt}"

    def enhanced_query(self, query: str) -> Dict:
        """增强的查询方法"""
        # 混合检索
        retrieval_results = self.hybrid_retrieve(query)

        # 生成专家响应
        response = self.generate_expert_response(query, retrieval_results)

        return {
            "response": response,
            "retrieval_stats": {
                "logs_count": len(retrieval_results["logs"]),
                "knowledge_count": len(retrieval_results["knowledge"]),
                "total_retrieved": len(retrieval_results["logs"]) + len(retrieval_results["knowledge"])
            },
            "top_logs": retrieval_results["logs"][:3],  # 返回前3个最相关的日志
            "top_knowledge": retrieval_results["knowledge"][:2]  # 返回前2个最相关的知识
        }


if __name__ == "__main__":
    # 初始化增强系统
    enhanced_system = EnhancedRAGSystem(
        log_path="./data/log",
        knowledge_path="./data/knowledge",
        llm="deepseek-r1:7b",
        embedding_model="bge-large:latest"
    )

    # 测试查询
    test_queries = [
        "数据库连接池耗尽如何解决？",
        "内存泄漏的常见原因和解决方案",
        "微服务认证失败怎么排查？"
    ]

    for query in test_queries:
        print(f"\n{'='*50}")
        print(f"查询: {query}")
        print(f"{'='*50}")

        result = enhanced_system.enhanced_query(query)
        print(f"检索统计: {result['retrieval_stats']}")
        print(f"响应: {result['response']}")
        print(f"\n相关日志: {len(result['top_logs'])} 条")
        print(f"相关知识: {len(result['top_knowledge'])} 条")